import numpy as np  # 提前导入 numpy
import requests
import json
from langchain_ollama import OllamaEmbeddings

# from llama_index.embeddings.huggingface import HuggingFaceEmbedding # 未使用，可注释或删除

# 配置ollama 服务
ollama_url = "http://127.0.0.1:11434"
OLLAMA_MODEL = "nomic-embed-text:latest"


def cosine_similarity(a, b):
    """计算余弦相似度，a 和 b 必须是 numpy 数组"""
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))


def get_ollama_model_metadata(model_name: str, base_url: str):
    """尝试从 Ollama API 获取模型元数据"""
    try:
        url = f"{base_url}/api/show"
        response = requests.post(url, json={"name": model_name})
        response.raise_for_status()

        data = response.text
        # Ollama show API 返回的是 Modelfile 文本，需要解析
        # 简单解析出参数量和大小（不完美，但能提供参考）

        details = {}
        for line in data.split('\n'):
            if line.startswith('parameter_count'):
                details['parameters_str'] = line.split()[1]
            elif line.startswith('model_size'):
                details['size_str'] = line.split()[1]

        return details

    except Exception as e:
        return {"error": f"无法连接或获取 Ollama 元数据: {e}"}


def test_embedding_consistency():
    # 打印模型详细信息 (通过 API 获取)
    model_details = get_ollama_model_metadata(OLLAMA_MODEL, ollama_url)

    print("=========================================")
    print(f"当前测试模型: {OLLAMA_MODEL}")
    if 'error' in model_details:
        print(f"⚠️ 警告: {model_details['error']}")
    else:
        print(f"⚙️ 模型参数量: {model_details.get('parameters_str', 'N/A')}")
        print(f"⚙️ 模型文件大小: {model_details.get('size_str', 'N/A')}")
    print("=========================================")

    try:
        # 初始化 Ollama 嵌入模型
        embed_model = OllamaEmbeddings(base_url=ollama_url, model=OLLAMA_MODEL)
    except Exception as e:
        print(f"❌ 初始化 OllamaEmbeddings 失败，请检查 Ollama 服务是否运行或模型是否已拉取: {e}")
        return

    text1 = "人工智能技术"
    text2 = "人工智能技术"  # 完全相同
    text3 = "人工 智能 技术"  # 稍有不同
    text4 = "机器学习算法"  # 完全不同

    print("\n=== 向量一致性测试 ===")

    texts = [text1, text2, text3, text4]
    embeddings = []

    # 循环生成向量
    for text in texts:
        embedding = embed_model.embed_query(text)
        embeddings.append(np.array(embedding))  # 将结果转换为 numpy 数组
        print(f"文本: '{text}'")
        print(f"向量前3维: {embedding[:3]}")

    # 获取向量维度
    if embeddings:
        dimension = embeddings[0].shape[0]
        print(f"⚙️ 向量维度: {dimension}")

    # --- 相似度分析 ---
    print(f"\n相似度分析 (余弦相似度):")

    # 完全相同文本：理想值 1.0
    sim_same = cosine_similarity(embeddings[0], embeddings[1])
    print(f"完全相同文本 (1 vs 2): {sim_same:.6f}")

    # 稍有不同：理想值接近 1.0
    sim_minor_diff = cosine_similarity(embeddings[0], embeddings[2])
    print(f"稍不同文本 (1 vs 3): {sim_minor_diff:.6f}")

    # 完全不同：理想值较低
    sim_diff = cosine_similarity(embeddings[0], embeddings[3])
    print(f"完全不同文本 (1 vs 4): {sim_diff:.6f}")


test_embedding_consistency()